import torch
import json
import time

from src.bert.my_model import BertClassifier, conf


def bert_predict(text_list):
    """
    对文本列表进行BERT预测

    Args:
        text_list: 文本列表

    Returns:
        dict: 符合指定格式的预测结果
    """
    start_time = time.time()

    # 确保输入是列表格式
    if isinstance(text_list, str):
        text_list = [text_list]

    model = BertClassifier()
    model.load_state_dict(torch.load(conf.save_model + 'bt_model.pth', map_location=conf.device))
    model.to(conf.device)
    model.eval()

    items = []

    with torch.no_grad():
        for text in text_list:
            item_start_time = time.time()

            # 记录输入文本
            input_text = text

            input_ids = conf.tokenizer.encode(text, return_tensors='pt', truncation=True, max_length=conf.max_len).to(conf.device)
            attention_mask = torch.ones(input_ids.shape, dtype=torch.long).to(conf.device)
            out_cat, out_label = model(input_ids, attention_mask)

            # 获取分类结果和概率
            predicted_cat_probs = torch.softmax(out_cat.data, dim=1)
            predicted_cat_confidence, predicted_cat = torch.max(predicted_cat_probs, 1)

            # 获取情感分析结果和概率
            label_prob = torch.sigmoid(out_label.data)

            # 读取categories.json文件
            with open(conf.class_path, 'r', encoding='utf-8') as f:
                categories = json.load(f)

            # 反转键值对，从id查找名称
            id_to_category = {v: k for k, v in categories.items()}

            # 确定情感倾向
            if label_prob.item() > 0.5:
                sentiment = 'positive'
                # score为预测为positive的置信度
                score = round(float(label_prob.item()), 2)
            else:
                sentiment = 'negative'
                # score为预测为negative的置信度，即1-预测为positive的概率
                score = round(float(1 - label_prob.item()), 2)

            # 获取标签（分类结果）
            tag = id_to_category.get(predicted_cat.item(), "未知分类")

            # 计算单条推理时间
            item_inference_time = (time.time() - item_start_time) * 1000

            # 构造单条结果
            item_result = {
                "text": input_text,
                "sentiment": sentiment,
                "score": score,
                "tags": [tag],
                "inferenceTime": round(item_inference_time, 2)
            }

            items.append(item_result)

        # 计算总推理时间
        total_inference_time = (time.time() - start_time) * 1000

        # 构造返回结果
        result = {
            'modelId': 'bert',
            'modelName': 'Bert模型',
            'total': len(text_list),
            'totalTime': round(total_inference_time, 2),
            'items': items
        }

        return result


if __name__ == '__main__':
    # 测试单个文本
    result1 = bert_predict('这手机真不错')
    print("单个文本预测结果:")
    print(result1)

    # 测试文本列表
    result2 = bert_predict(['这手机真不错', '质量不怎么样。京东的东西越来越不行了,'])
    print("\n多个文本预测结果:")
    print(result2)
